import cv2
import numpy as np
import pydensecrf.densecrf as dcrf

from cv2 import imread, imwrite

from pydensecrf.utils import unary_from_labels, create_pairwise_bilateral, create_pairwise_gaussian


def CRF_point(img, labels, has_unsure):
    # 计算predicted_image中的类数。
    n_labels = len(set(labels.flat))
    #  n_labels = len(set(labels.flat)) - int(HAS_UNK) ##如果有不确定区域，用这一行代码替换上一行

    # 使用densecrf2d类
    d = dcrf.DenseCRF2D(img.shape[1], img.shape[0], n_labels)

    # 得到一元势（负对数概率）
    U = unary_from_labels(labels, n_labels, gt_prob=0.7, zero_unsure=has_unsure)
    # U = unary_from_labels(labels, n_labels, gt_prob=0.2, zero_unsure=HAS_UNK)## 如果有不确定区域，用这一行代码替换上一行
    d.setUnaryEnergy(U)

    # 增加了与颜色无关的术语，只是位置-----会惩罚空间上孤立的小块分割,即强制执行空间上更一致的分割
    d.addPairwiseGaussian(sxy=(3, 3), compat=5, kernel=dcrf.DIAG_KERNEL,
                          normalization=dcrf.NORMALIZE_SYMMETRIC)

    # 增加了颜色相关术语，即特征是(x,y,r,g,b)-----使用局部颜色特征来细化它们
    d.addPairwiseBilateral(sxy=(80, 80), srgb=(13, 13, 13), rgbim=img, compat=10,
                           kernel=dcrf.DIAG_KERNEL,
                           normalization=dcrf.NORMALIZE_SYMMETRIC)

    ####################################
    ###         做推理和计算         ###
    ####################################

    Q = d.inference(10)
    # 找出每个像素最可能的类
    MAP = np.argmax(Q, axis=0)

    return MAP

if __name__  =="__main__":
    #img = imread("data/input.png")
    img = imread("E:\\SRTP\\DATA\\ld-image\\ld-image_Panoramic_000109_11651_123-4354.jpg")
    img = cv2.resize(img,(2048,1024))
    anno_bgr = imread("./data/proj.png")
    anno_lbl = anno_bgr[:, :, 0] + (anno_bgr[:, :, 1] << 8) + (anno_bgr[:, :, 2] << 16)
    # 将uint32颜色转换为1,2,...
    colors, labels = np.unique(anno_lbl, return_inverse=True)

    # 创建从predicted_image到32位整数颜色的映射。
    colorize = np.empty((len(colors), 3), np.uint8)
    colorize[:, 0] = (colors & 0x0000FF)
    colorize[:, 1] = (colors & 0x00FF00) >> 8
    colorize[:, 2] = (colors & 0xFF0000) >> 16
    res = CRF_point(img, labels,False)
    MAP = colorize[res, :]
    imwrite("./data/output1.jpg",MAP.reshape(img.shape))